package com.hhq.sso.protocol.oauth2.config;

import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.text.ParseException;
import java.util.UUID;
import java.util.function.Function;

import com.hhq.sso.protocol.oauth2.constants.RedisConstants;
import com.nimbusds.jose.jwk.source.ImmutableJWKSet;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Import;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.security.config.Customizer;
import org.springframework.security.config.annotation.web.builders.HttpSecurity;
import org.springframework.security.config.annotation.web.configuration.EnableWebSecurity;
import org.springframework.security.config.annotation.web.configurers.AbstractHttpConfigurer;
import org.springframework.security.core.userdetails.UserDetailsService;
import org.springframework.security.crypto.bcrypt.BCryptPasswordEncoder;
import org.springframework.security.crypto.password.PasswordEncoder;
import org.springframework.security.oauth2.core.oidc.OidcUserInfo;
import org.springframework.security.oauth2.core.oidc.endpoint.OidcParameterNames;
import org.springframework.security.oauth2.jwt.JwtDecoder;
import org.springframework.security.oauth2.jwt.JwtEncoder;
import org.springframework.security.oauth2.jwt.NimbusJwtEncoder;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationConsentService;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.config.annotation.web.configuration.OAuth2AuthorizationServerConfiguration;
import org.springframework.security.oauth2.server.authorization.config.annotation.web.configurers.OAuth2AuthorizationServerConfigurer;
import org.springframework.security.oauth2.server.authorization.oidc.authentication.OidcUserInfoAuthenticationContext;
import org.springframework.security.oauth2.server.authorization.oidc.authentication.OidcUserInfoAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.settings.AuthorizationServerSettings;
import org.springframework.security.oauth2.server.authorization.token.*;
import org.springframework.security.oauth2.server.resource.authentication.JwtAuthenticationToken;
import org.springframework.security.web.SecurityFilterChain;

import com.hhq.sso.protocol.oauth2.entity.SsoUser;
import com.hhq.sso.protocol.oauth2.repository.AuthorizationConsentRepository;
import com.hhq.sso.protocol.oauth2.repository.AuthorizationRepository;
import com.hhq.sso.protocol.oauth2.repository.JpaRegisteredClientRepository;
import com.hhq.sso.protocol.oauth2.repository.SsoUserRepository;
import com.hhq.sso.protocol.oauth2.service.JpaOAuth2AuthorizationConsentService;
import com.hhq.sso.protocol.oauth2.service.JpaOAuth2AuthorizationService;
import com.hhq.sso.protocol.oauth2.service.SsoUserService;
import com.nimbusds.jose.jwk.JWKSet;
import com.nimbusds.jose.jwk.RSAKey;
import com.nimbusds.jose.jwk.source.JWKSource;
import com.nimbusds.jose.proc.SecurityContext;

import lombok.extern.slf4j.Slf4j;
import org.springframework.util.ObjectUtils;

/**
 * <p>
 * Description: 授权服务器配置
 * </p>
 *
 * @author : xiaodong.yang
 * @date : 2024/6/14 17:04
 */
@Configuration
@Import(OAuth2AuthorizationServerConfiguration.class)
@EnableWebSecurity
@Slf4j
public class AuthorizationServerConfig {

    @Autowired
    private JpaRegisteredClientRepository jpaRegisteredClientRepository;

    @Autowired
    private AuthorizationConsentRepository authorizationConsentRepository;

    @Autowired
    private AuthorizationRepository authorizationRepository;
    @Autowired
    private SsoUserRepository ssoUserRepository;

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    @Bean
    public JWKSource<SecurityContext> jwkSource() {
        String jwkSetCache = (String)redisTemplate.opsForValue().get(RedisConstants.AUTHORIZATION_JWS_PREFIX_KEY);
        if (ObjectUtils.isEmpty(jwkSetCache)) {
            KeyPair keyPair = generateRsaKey();
            RSAPublicKey publicKey = (RSAPublicKey)keyPair.getPublic();
            RSAPrivateKey privateKey = (RSAPrivateKey)keyPair.getPrivate();
            RSAKey rsaKey =
                new RSAKey.Builder(publicKey).privateKey(privateKey).keyID(UUID.randomUUID().toString()).build();
            JWKSet jwkSet = new JWKSet(rsaKey);
            // 转为json字符串
            String jwkSetString = jwkSet.toString(Boolean.FALSE);
            redisTemplate.opsForValue().set(RedisConstants.AUTHORIZATION_JWS_PREFIX_KEY, jwkSetString);
            return new ImmutableJWKSet<>(jwkSet);
        }
        // 解析存储的jws
        JWKSet jwkSet;
        try {
            jwkSet = JWKSet.parse(jwkSetCache);
            return new ImmutableJWKSet<>(jwkSet);
        } catch (ParseException e) {
            throw new RuntimeException(e);
        }

    }

    /**
     * 配置token生成器
     *
     * @return
     */
    @Bean
    public OAuth2TokenGenerator<?> tokenGenerator() {
        JwtEncoder jwtEncoder = new NimbusJwtEncoder(jwkSource());
        JwtGenerator jwtGenerator = new JwtGenerator(jwtEncoder);
        jwtGenerator.setJwtCustomizer(jwtCustomizer());
        OAuth2AccessTokenGenerator accessTokenGenerator = new OAuth2AccessTokenGenerator();
        accessTokenGenerator.setAccessTokenCustomizer(accessTokenCustomizer());
        OAuth2RefreshTokenGenerator refreshTokenGenerator = new OAuth2RefreshTokenGenerator();
        return new DelegatingOAuth2TokenGenerator(jwtGenerator, accessTokenGenerator, refreshTokenGenerator);
    }

    /**
     * jwt定制化
     *
     * @return
     */
    @Bean
    public OAuth2TokenCustomizer<JwtEncodingContext> jwtCustomizer() {
        return context -> {
            // 从acessToken获取用户信息
            if (context.getTokenType().equals(OAuth2TokenType.ACCESS_TOKEN)) {
                // Customize headers/claims for access_token
                SsoUser userInfo = (SsoUser)userDetailsService().loadUserByUsername(context.getPrincipal().getName());
                context.getClaims().claims(claims -> claims.putAll(userInfo.info()));
            } else if (context.getTokenType().getValue().equals(OidcParameterNames.ID_TOKEN)) {
                // Customize headers/claims for id_token
            }
        };
    }

    /**
     * accessToken定制化
     *
     * @return
     */
    @Bean
    public OAuth2TokenCustomizer<OAuth2TokenClaimsContext> accessTokenCustomizer() {
        return context -> {
            OAuth2TokenClaimsSet.Builder claims = context.getClaims();
            // Customize claims
        };
    }

    @Bean
    public OAuth2AuthorizationService authorizationService() {
        return new JpaOAuth2AuthorizationService(authorizationRepository, jpaRegisteredClientRepository);
    }

    @Bean
    public OAuth2AuthorizationConsentService getOAuth2AuthorizationConsentService() {
        return new JpaOAuth2AuthorizationConsentService(authorizationConsentRepository, jpaRegisteredClientRepository);
    }

    /**
     * @param http
     * @return
     * @throws Exception
     */
    @Bean
    public SecurityFilterChain authorizationServerSecurityFilterChain(HttpSecurity http) throws Exception {
        OAuth2AuthorizationServerConfigurer authorizationServerConfigurer = new OAuth2AuthorizationServerConfigurer();
        http.apply(authorizationServerConfigurer);
        http
            // 禁用匿名身份验证
            .anonymous(AbstractHttpConfigurer::disable);
        Function<OidcUserInfoAuthenticationContext, OidcUserInfo> userInfoMapper = (context) -> {
            OidcUserInfoAuthenticationToken authentication = context.getAuthentication();
            JwtAuthenticationToken principal = (JwtAuthenticationToken)authentication.getPrincipal();
            return new OidcUserInfo(principal.getToken().getClaims());
        };
        // 跨域
        authorizationServerConfigurer.registeredClientRepository(jpaRegisteredClientRepository)
            .authorizationService(authorizationService())
            .authorizationConsentService(getOAuth2AuthorizationConsentService())
            .authorizationServerSettings(authorizationServerSettings()).tokenGenerator(tokenGenerator())
            .clientAuthentication(
                clientAuthentication -> clientAuthentication.errorResponseHandler((request, response, e) -> {
                    log.error("clientAuthentication error", e);
                    response.setContentType(MediaType.APPLICATION_JSON_UTF8_VALUE);
                    response.setStatus(HttpStatus.UNAUTHORIZED.value());
                    response.getWriter().write(e.getMessage());
                }))
            .authorizationEndpoint(
                // authorize 授权异常处理
                authorizationEndpoint -> authorizationEndpoint.errorResponseHandler((request, response, e) -> {
                    log.error("authorizationEndpoint error", e);
                    response.setContentType(MediaType.APPLICATION_JSON_UTF8_VALUE);
                    response.setStatus(HttpStatus.UNAUTHORIZED.value());
                    response.getWriter().write("没有认证");
                }))
            .deviceAuthorizationEndpoint(deviceAuthorizationEndpoint -> {
            }).deviceVerificationEndpoint(deviceVerificationEndpoint -> {
            }).tokenEndpoint(tokenEndpoint -> tokenEndpoint
                // token异常返回处理
                .errorResponseHandler((request, response, e) -> {
                    log.error("tokenEndpoint error", e);
                    response.setContentType(MediaType.APPLICATION_JSON_UTF8_VALUE);
                    response.setStatus(HttpStatus.UNAUTHORIZED.value());
                    response.getWriter().write(e.getMessage());
                }))
            .tokenIntrospectionEndpoint(tokenIntrospectionEndpoint -> {
            }).tokenRevocationEndpoint(tokenRevocationEndpoint -> {
            }).authorizationServerMetadataEndpoint(authorizationServerMetadataEndpoint -> {
            }).oidc((oidc) -> oidc.userInfoEndpoint((userInfo) -> userInfo.userInfoMapper(userInfoMapper)));

        // 配置资源服务器，通过acess_token获取用户信息
        http.oauth2ResourceServer((resourceServer) -> resourceServer.jwt(Customizer.withDefaults()));

        http.authorizeHttpRequests((authorizeHttpRequests) -> authorizeHttpRequests
            // 放行的接口
            // .requestMatchers("/oauth2/**").permitAll()
            .requestMatchers("/assets/**", "/login", "/captcha").permitAll()
            // 其他接口全部需要鉴权
            .anyRequest().authenticated());

        http
            // 关闭跨域
            .csrf(AbstractHttpConfigurer::disable).formLogin(formLogin -> formLogin.loginPage("/login"));
        // 图形验证码
        // http.addFilterBefore(new CaptchaAuthenticationFilter("/login", CommonConstant.CAPTCHA_SESSION_KEY),
        // UsernamePasswordAuthenticationFilter.class);
        return http.build();
    }

    /**
     * 配置用户数据库查询
     *
     * @return
     */
    @Bean
    public UserDetailsService userDetailsService() {
        return new SsoUserService(ssoUserRepository);
    }

    @Bean
    public PasswordEncoder passwordEncoder() {
        return new BCryptPasswordEncoder();
    }

    /**
     * 自定义服务器配置
     *
     * @return
     */
    @Bean
    public AuthorizationServerSettings authorizationServerSettings() {
        return AuthorizationServerSettings.builder().authorizationEndpoint("/oauth2/authorize")
            .deviceAuthorizationEndpoint("/oauth2/device_authorization")
            .deviceVerificationEndpoint("/oauth2/device_verification").tokenEndpoint("/oauth2/token")
            .tokenIntrospectionEndpoint("/oauth2/introspect").tokenRevocationEndpoint("/oauth2/revoke")
            .jwkSetEndpoint("/oauth2/jwks").oidcLogoutEndpoint("/connect/logout")
            .oidcUserInfoEndpoint("/connect/userinfo").oidcClientRegistrationEndpoint("/connect/register").build();
    }

    @Bean
    public JwtDecoder jwtDecoder(JWKSource<SecurityContext> jwkSource) {
        return OAuth2AuthorizationServerConfiguration.jwtDecoder(jwkSource);
    }

    private static KeyPair generateRsaKey() {
        KeyPair keyPair;
        try {
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
            keyPairGenerator.initialize(2048);
            keyPair = keyPairGenerator.generateKeyPair();
        } catch (Exception ex) {
            throw new IllegalStateException(ex);
        }
        return keyPair;
    }
}
